{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 提示词 #\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI(\n",
    "    #api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    # api_key=\"sk-b571bfbe652b4ec68ac0491e33949622\", # 这种写法不好，泄露了api-key。 回头正式部署时改掉。\n",
    "    # base_url=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    api_key=\"ollama\",\n",
    "    base_url=\"http://192.168.20.43:11434/v1\"\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_qa_prompt(query, context_str):\n",
    "    prompt = f\"\"\"\\\n",
    "    ---------------------\\\n",
    "    # OBJECTIVE #\n",
    "    You are a QA assistant.\n",
    "    <Context> is below. Given the context information and not prior knowledge, \\\n",
    "    answer the query.\\\n",
    "    ---------------------\\\n",
    "    # STYLE #\n",
    "    风格应保持结构清晰且简洁易读。\n",
    "    ---------------------\\\n",
    "    Reply in Chinese. \\\n",
    "    ---------------------\\\n",
    "    # Context #\n",
    "    {context_str}\\\n",
    "    ---------------------\\\n",
    "    \"\"\"\n",
    "    completion = client.chat.completions.create(\n",
    "        model=\"qwen2-7b-instruct\",\n",
    "        messages=[{'role': 'system', 'content': prompt},\n",
    "                  {'role': 'user', 'content': query}],\n",
    "        )\n",
    "    return completion.choices[0].message.content\n",
    "\n",
    "#测试用例\n",
    "context_str = \"\"\"\n",
    "    我们拥有每家加油站每小时的净利润数据，你可以通过API获取，get_profit(GasStation_ID, time_period)\n",
    "    我们也拥有全部加油站每小时的净利润数据，你可以通过API获取，get_TotalProfit(time_period)\n",
    "\"\"\"\n",
    "query1 = \"\"\"\n",
    "我们哪家加油站的营业时间需要调整么？\n",
    "\"\"\"   \n",
    "summary_text = llm_qa_prompt(query1,context_str) \n",
    "print(summary_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm_check_context(query, context_str):\n",
    "    prompt = \"\"\n",
    "\n",
    "    query_new = f\"\"\"\n",
    "    ---------------------\\\n",
    "    # OBJECTIVE #\\\n",
    "    你的任务是判断<context>是否提供了足够的信息来回答<query>。请遵循以下严格的判断步骤：\\\n",
    "    1.仔细阅读<context>，禁止对<context>的内容做任何理解、分析、推断和延展；禁止对语境或场景进行推断和延展。\n",
    "    2.仅思考<context>和<query>之间的相关性。\n",
    "    若<context>中未提供具体细节来回答<query>，请回复\"False\"；若<context>中无须推断就包括了具体细节来回答<query>，请回复\"True\"。并给出不超过10个字的理由。\\\n",
    "    ----------------\\\n",
    "    <context>: 纽约是美国的大城市。\\\n",
    "    <query>: 请告诉我中国的全称？\\\n",
    "    <ai>: False，未提供关于中国全称的相关信息。\n",
    "    \n",
    "    <context>: 猫是胎生动物，鱼是卵生动物。猫和鱼不是同一种类的动物。\\\n",
    "    <query>: 猫和鱼是同一种类的动物么？\\\n",
    "    <ai>: True，已经明确指出“猫和鱼不是同一种类的动物”，这直接回\\答了查询“猫和鱼是同一种类的动物么？\\\n",
    "\n",
    "    <context>:{context_str}\\\n",
    "    <query>: {query}\\\n",
    "    <ai>: \n",
    "    \"\"\"\n",
    "    completion = client.chat.completions.create(\n",
    "        model=\"qwen2-7b-instruct\",\n",
    "        messages=[{'role': 'system', 'content': prompt},\n",
    "                  {'role': 'user', 'content': query_new}],\n",
    "        )\n",
    "    return completion.choices[0].message.content\n",
    "\n",
    "\n",
    "# 测试用例\n",
    "context_str1 = \"\"\"\n",
    "section 1 start:石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对石油工作者很重要 sction 1 ends.\n",
    "\"\"\"\n",
    "context_str2 = \"\"\"\n",
    "《油气知识概述》\n",
    "一、引言\n",
    "石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对于认识能源格局、推动可持续发展以及保障国家能源安全都具有重大意义。\n",
    "\"\"\"\n",
    "context_str3 = \"\"\"\n",
    "section 1 start:《油气知识概述》\n",
    "一、引言\n",
    "石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对于认识能源格局、推动可持续发展以及保障国家能源安全都具有重大意义。\n",
    "section 1 ends.\n",
    "\"\"\"\n",
    "context_str4 = \"油气知识\"\n",
    "query = \"天然气的开采方法\"   \n",
    "bool_contextIsEnough = llm_check_context(query,context_str3) \n",
    "print(bool_contextIsEnough)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 两阶段RAG回复\n",
    "import re\n",
    "\n",
    "def get_trueOrFalse(check_context):\n",
    "        # 使用strip方法移除字符串首尾的空格\n",
    "    trimmed_text = check_context.strip()\n",
    "    print(trimmed_text)\n",
    "    # 检查清理后的字符串是否只包含True或False\n",
    "    if trimmed_text in ['True', 'False']:\n",
    "        # 如果是，直接输出\n",
    "        print(trimmed_text)\n",
    "        return trimmed_text\n",
    "    else:\n",
    "        # 如果不是，使用split方法以逗号分隔字符串，并取第一部分\n",
    "        first_part = re.split(r'[，,]', trimmed_text)[0]\n",
    "        print(first_part)\n",
    "        return first_part\n",
    "\n",
    "def llm_qa_2stage(query, context):\n",
    "    text_answer = \"\"\n",
    "    #阶段一：用llm判断是否给出的context包括的信息是否足够回答query\n",
    "    check_context = llm_check_context(query,context)\n",
    "    TrueOrFalse = get_trueOrFalse(check_context)\n",
    "\n",
    "    #阶段二：如果足够回答query，再调用llm回答\n",
    "    if TrueOrFalse == \"True\":\n",
    "        text_answer = llm_qa_prompt(query,context) \n",
    "        return text_answer\n",
    "    elif TrueOrFalse == \"False\":\n",
    "        return \"这个问题不在知识范围之内，现阶段暂时无法回答。\"\n",
    "    else:\n",
    "        return \"llm_qa_2stage检测失败\"\n",
    "\n",
    "# 测试用例\n",
    "context_str1 = \"\"\"\n",
    "section 1 start:石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对石油工作者很重要 sction 1 ends.\n",
    "\"\"\"\n",
    "context_str2 = \"\"\"\n",
    "《油气知识概述》\n",
    "一、引言\n",
    "石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对于认识能源格局、推动可持续发展以及保障国家能源安全都具有重大意义。\n",
    "\"\"\"\n",
    "context_str3 = \"\"\"\n",
    "section 1 start:《油气知识概述》\n",
    "一、引言\n",
    "石油和天然气作为重要的能源资源，在全球经济和人们的日常生活中发挥着至关重要的作用。从驱动汽车、供暖到为工业生产提供动力，油气的身影无处不在。了解油气知识，对于认识能源格局、推动可持续发展以及保障国家能源安全都具有重大意义。\n",
    "section 1 ends.\n",
    "\"\"\"\n",
    "context_str4 = \"\"\"\n",
    "油气知识\n",
    "\"\"\"\n",
    "query = \"天然气的开采方法\"   \n",
    "llm_output = llm_qa_2stage(query,context_str3) \n",
    "print(llm_output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
